#!/bin/bash

# author: fangfang
# created: 2022年12月6日
# updated:
# desc: 该脚本用于AdsorptionEnergyPrediction推理作业运行脚本
#       所用数据集路径为启pod时默认挂的oss路径

# 该脚本接收3个参数
# --input_dir 训练数据集所在文件夹，若不指定则默认使用已挂载数据集
# --output_dir 结果输出目录
# --gpus 该作业使用的gpu卡数


for key in "$@"; do
  case $key in
  --input_dir=*)
    input_dir="${key#*=}"
    shift # past argument=value
    ;;
  --output_dir=*)
    output_dir="${key#*=}"
    shift # past argument=value
    ;;
  --gpus=*)
    gpus="${key#*=}"
    shift # past argument=value
    ;;
  --default)
    DEFAULT=YES
    shift # past argument with no value
    ;;
  *) ;;

  esac
done

input_dir=${input_dir%*/}
files=$(ls $input_dir)

model_path="/dros/common/ecosystem/MEDICINE_COMPUTING/apps/AdsorptionEnergyPrediction/MODEL/best_checkpoint.pt"

for file in $files; do
  echo "start: ${file}"
  data="$input_dir/$file"
  /ocp22/bin/python -u -m torch.distributed.launch --nproc_per_node=$gpus main.py --distributed --num-gpus $gpus --mode predict --config-yml ./configs/is2re/all/scn_proj/SCN-6-12-t1-b1-2M_proj.yml --checkpoint $model_path --test-dataset $data
done

mv results $output_dir
